from langchain_community.vectorstores import Chroma
from langchain_openai import AzureChatOpenAI, AzureOpenAIEmbeddings
from langchain.prompts import PromptTemplate
from docx import Document
import os
import datetime

# ==== Azure API 配置 ====
os.environ["AZURE_OPENAI_API_KEY"] = "EVlHLUCSdiDb0yVKwkm6USjrqshc2hZxyoLD5EI7PBzl3coU7ApTJQQJ99BCACYeBjFXJ3w3AAABACOGVmn4"
os.environ["AZURE_OPENAI_ENDPOINT"] = "https://eastus-0303.openai.azure.com/"
os.environ["AZURE_OPENAI_API_VERSION"] = "2024-08-01-preview"

# ===== 1. 读取向量数据库 =====
embeddings = AzureOpenAIEmbeddings(
    deployment="text-embedding-3-large-1",  # 你部署的 embedding 模型
    model="text-embedding-3-large",
    api_key=os.environ["AZURE_OPENAI_API_KEY"],
    azure_endpoint=os.environ["AZURE_OPENAI_ENDPOINT"],
    api_version="2023-05-15"
)

db = Chroma(persist_directory="company_db", embedding_function=embeddings)

# ===== 2. 初始化 Azure GPT-4o-1 =====
llm = AzureChatOpenAI(
    deployment_name="gpt-4o-1",  # 你 Azure 部署的对话模型名
    model="gpt-4o",
    api_key=os.environ["AZURE_OPENAI_API_KEY"],
    azure_endpoint=os.environ["AZURE_OPENAI_ENDPOINT"],
    api_version=os.environ["AZURE_OPENAI_API_VERSION"],
    temperature=0
)

# ===== 3. Prompt 模板 =====
template = """你是一名经验丰富的售前技术专家。  
你的任务是基于以下资料，帮助销售人员快速理解客户的需求和痛点，并结合公司现有的产品、服务对象，指出可能的合作点和面。  

请遵循以下规则：  
1. 回答时要结构化，包含：
   - 客户需求/痛点  
   - 合作契机
   - 智用与合作公司的异同点
   - 公司能够提供的价值或解决方案  
   - 潜在合作点和面  
   - 下一步建议  
2. 如果资料中没有相关信息，请明确说明“未找到相关资料”，不要编造内容。  
3. 输出时要简洁清晰，便于销售快速把握重点。  

资料内容：  
{context}

用户问题：  
{question}

请基于以上信息给出专业回答：
"""
prompt = PromptTemplate(
    input_variables=["context", "question"],
    template=template
)

# ===== 4. 创建 Word 文档（文件名带日期）=====
doc_filename = f"rag_qa_{datetime.datetime.now().strftime('%Y%m%d_%H%M%S')}.docx"
doc = Document()
doc.add_heading("RAG 问答记录", level=1)

# ===== 5. 运行问答 =====
while True:
    question = input("\n请输入你的问题（输入 'exit' 退出）：")
    if question.lower() == "exit":
        print(f"✅ 已保存到 {doc_filename}")
        break

    # 从数据库检索
    docs = db.similarity_search(question, k=3)
    context = "\n\n".join([doc.page_content for doc in docs])

    # 构造 prompt
    final_prompt = prompt.format(context=context, question=question)

    # 调用 GPT
    response = llm.invoke(final_prompt)
    answer = response.content
    print("\n💡 回答：", answer)

    # === 保存到 Word ===
    doc.add_heading(f"问题：{question}", level=2)
    doc.add_paragraph(answer)

    # 每次问答后立即保存一次
    doc.save(doc_filename)
    print(f"📄 已写入并保存到 {doc_filename}")
